package ink.mint.convolution;

import ink.mint.activation.Activateable;
import ink.mint.activation.ActivationFunction;
import ink.mint.activation.ActivationType;
import ink.mint.matrix.MatrixExtensions;

import java.io.Serializable;

public class ConvolutionNeural implements Serializable {

    MatrixExtensions matrix;
    Activateable activateable;
    ActivationType activationType;
    PoolingType poolingType;
    double[][] input;
    double[][] kernel;
    double bias;
    transient double[][] featureMap;
    transient double[][] activateMap;
    transient double[][] poolingMap;
    transient double[][] poolingMapDelta;
    transient double[][] upSampleDelta;
    transient double[][] derivationActivateMap;
    transient double[][] kernelDelta;
    transient double biasDelta;
    transient double[][] derivationActivateMapDelta;
    int kernelStride;
    int kernelSize;
    int convolutionPadding;

    int poolingStride;
    int poolingSize;
    int poolingPadding;

    int[][] poolingMaxIndex;

    public ConvolutionNeural() {
        matrix = new MatrixExtensions();
        activateable = new ActivationFunction();
    }

    public ConvolutionNeural input(double[][] input) {
        this.input = input;
        return this;
    }

    private void setKernelStride(int kernelStride) {
        this.kernelStride = kernelStride;
    }

    private void setKernelSize(int kernelSize) {
        this.kernelSize = kernelSize;
    }

    private void setConvolutionPadding(int convolutionPadding) {
        this.convolutionPadding = convolutionPadding;
    }

    private void setActivationType(ActivationType activationType) {
        this.activationType = activationType;
    }

    private void setPoolingSize(int poolingSize) {
        this.poolingSize = poolingSize;
    }

    private void setPoolingStride(int poolingStride) {
        this.poolingStride = poolingStride;
    }

    private void setPoolingPadding(int poolingPadding) {
        this.poolingPadding = poolingPadding;
    }

    private void setPoolingType(PoolingType poolingType) {
        this.poolingType = poolingType;
    }

    public ConvolutionNeural setConvolutionParameter(int kernelSize, int kernelStride, int convolutionPadding, ActivationType activationType) {
        setKernelSize(kernelSize);
        setKernelStride(kernelStride);
        setConvolutionPadding(convolutionPadding);
        setActivationType(activationType);
        return this;
    }

    public ConvolutionNeural setPoolingParameter(int poolingSize, int poolingStride, int poolingPadding, PoolingType poolingType) {
        setPoolingSize(poolingSize);
        setPoolingStride(poolingStride);
        setPoolingPadding(poolingPadding);
        setPoolingType(poolingType);
        return this;
    }

//    private void initKernelAndBias() {
//        this.kernel = matrix.createRandomKernel(kernelSize);
//        this.bias = Math.random() - 0.5;
//    }

    // ConvolutionNeural.java 修改初始化条件
    private void initKernelAndBias() {
        // 添加null检查和维度验证
        if (kernel == null ||
                kernel.length != kernelSize ||
                kernel[0].length != kernelSize) {

            this.kernel = matrix.createRandomKernel(kernelSize);
            this.bias = Math.random() - 0.5;
        }
    }


    private void updataKernelAndBias(double[][] kernel, double bias) {
        this.kernel = kernel;
        this.bias = bias;
    }

    public double[][] convolution(double[][] kernel, double bias,boolean isTrain) {
        int ir = input.length;
        int ic = input[0].length;
        int kr = kernelSize;
        int kc = kernelSize;
//        int fr = (ir - kr + convolutionPadding) / kernelStride + 1;
//        int fc = (ic - kc + convolutionPadding) / kernelStride + 1;

// ConvolutionNeural.java 中卷积输出尺寸计算
//        int fr = (ir - kr + convolutionPadding) / kernelStride + 1; // ❌错误公式
// 正确公式应包含两侧padding：
        int fr = (ir + 2*convolutionPadding - kr) / kernelStride + 1; // ✅
        int fc = (ic + 2*convolutionPadding - kc) / kernelStride + 1;
        double[][] featureMap = new double[fr][fc];

        if(isTrain) {
            if (kernel.length != kernelSize) {
                initKernelAndBias();
            } else {
                updataKernelAndBias(kernel, bias);
            }
        }
        featureMap = matrix.featureMap(input, this.kernel, this.bias, kernelStride, convolutionPadding);
        this.featureMap = featureMap;
        return featureMap;
    }

    public double[][] activate() {
        double[][] activateMap = activateable.ActivationSolve(featureMap, activationType);
        this.activateMap = activateMap;

        return activateMap;
    }

    public double[][] pooling() {
        int ar = activateMap.length;
        int ac = activateMap[0].length;
//        int pr = (ar - poolingSize + poolingPadding) / poolingStride + 1;
//        int pc = (ac - poolingSize + poolingPadding) / poolingStride + 1;

        int pr = (ar + 2 * poolingPadding - poolingSize) / poolingStride + 1;
        int pc = (ac + 2 * poolingPadding - poolingSize) / poolingStride + 1;


        double[][] poolingMap = new double[pr][pc];
        switch (poolingType) {
            case MaxPooling:
                poolingMap = matrix.maxPooling(activateMap, poolingSize, poolingStride, poolingPadding);
                poolingMaxIndex = matrix.getMaxIndex();
                this.poolingMap = poolingMap;
                return poolingMap;
            case AveragePooling:
                poolingMap = matrix.averagePooling(activateMap, poolingSize, poolingStride, poolingPadding);
                this.poolingMap = poolingMap;
                return poolingMap;
            default:
                return poolingMap;
        }
    }


    public MatrixExtensions getMatrix() {
        return matrix;
    }

    public void setMatrix(MatrixExtensions matrix) {
        this.matrix = matrix;
    }

    public Activateable getActivateable() {
        return activateable;
    }

    public void setActivateable(Activateable activateable) {
        this.activateable = activateable;
    }

    public ActivationType getActivationType() {
        return activationType;
    }

    public PoolingType getPoolingType() {
        return poolingType;
    }

    public double[][] getInput() {
        return input;
    }

    public void setInput(double[][] input) {
        this.input = input;
    }

    public double[][] getKernel() {
        return kernel;
    }

    public void setKernel(double[][] kernel) {
        this.kernel = kernel;
    }

    public double getBias() {
        return bias;
    }

    public void setBias(double bias) {
        this.bias = bias;
    }

    public double[][] getFeatureMap() {
        return featureMap;
    }

    public void setFeatureMap(double[][] featureMap) {
        this.featureMap = featureMap;
    }

    public double[][] getActivateMap() {
        return activateMap;
    }

    public void setActivateMap(double[][] activateMap) {
        this.activateMap = activateMap;
    }

    public double[][] getPoolingMap() {
        return poolingMap;
    }

    public void setPoolingMap(double[][] poolingMap) {
        this.poolingMap = poolingMap;
    }

    public int getKernelStride() {
        return kernelStride;
    }

    public int getKernelSize() {
        return kernelSize;
    }

    public int getConvolutionPadding() {
        return convolutionPadding;
    }

    public int getPoolingStride() {
        return poolingStride;
    }

    public int getPoolingSize() {
        return poolingSize;
    }

    public int getPoolingPadding() {
        return poolingPadding;
    }

    public int[][] getPoolingMaxIndex() {
        return poolingMaxIndex;
    }

    public void setPoolingMaxIndex(int[][] poolingMaxIndex) {
        this.poolingMaxIndex = poolingMaxIndex;
    }

    public double[][] getPoolingMapDelta() {
        return poolingMapDelta;
    }

    public ConvolutionNeural setPoolingMapDelta(double[][] poolingMapDelta) {
        this.poolingMapDelta = poolingMapDelta;
        return this;
    }

    public double[][] getUpSampleDelta() {
        return upSampleDelta;
    }

    public void setUpSampleDelta(double[][] upSampleDelta) {
        this.upSampleDelta = upSampleDelta;
    }

    public double[][] getDerivationActivateMap() {
        return derivationActivateMap;
    }

    public void setDerivationActivateMap(double[][] derivationActivateMap) {
        this.derivationActivateMap = derivationActivateMap;
    }

    public double[][] getKernelDelta() {
        return kernelDelta;
    }

    public void setKernelDelta(double[][] kernelDelta) {
        this.kernelDelta = kernelDelta;
    }

    public double[][] getDerivationActivateMapDelta() {
        return derivationActivateMapDelta;
    }

    public void setDerivationActivateMapDelta(double[][] derivationActivateMapDelta) {
        this.derivationActivateMapDelta = derivationActivateMapDelta;
    }

    public double getBiasDelta() {
        return biasDelta;
    }

    public void setBiasDelta(double biasDelta) {
        this.biasDelta = biasDelta;
    }


    public void setInputDelta(double[][] inputDelta) {
    }
}
